#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time   : 2023-09-19 21:30
# @Author : hard work
# @File    : check_pcap_result_util.py
# @Description :
import os
from pyshark.capture.file_capture import FileCapture
from loguru import logger


# tshark_path =E:\Program Files\Wireshark\tshark.exe

class Check_pcap_res_base(object):
    def __init__(self, pcap_file_path):
        self.__file_name__ = pcap_file_path
        self.__wireshark_file__ = r"E:\\Program Files\\Wireshark"
        self.__PKT_LST__ = []
        self.__parse_pkt__()

    def __parse_pkt__(self):
        if os.path.isfile(self.__file_name__):
            try:
                # pcap = pyshark.FileCapture(self.__file_name__, tshark_path=self.__wireshark_file__ + r"\\tshark.exe")
                # pcap = pyshark.FileCapture(self.__file_name__)
                pcap = FileCapture(self.__file_name__, tshark_path=self.__wireshark_file__ + r"\\tshark.exe")

                for pkt in pcap:
                    self.__PKT_LST__.append(pkt)
                print("Info ", "%s解析完毕" % self.__file_name__)
                pcap.close()
                return self.__PKT_LST__
            except RuntimeError as e:
                print("pyshark运行时发生错误", e)
            except FileExistsError as e:
                print("文件不存在", e)
        else:
            return

    def get_src_mac(self):
        mac_sa_lst = []
        for pkt in self.__PKT_LST__:
            mac_sa = str(pkt.eth.src)
            mac_sa_lst.append(mac_sa)
        return mac_sa_lst

    def get_dst_mac(self):
        mac_da_lst = []
        for pkt in self.__PKT_LST__:
            mac_da = str(pkt.eth.dst)
            mac_da_lst.append(mac_da)
        return mac_da_lst

    def get_vid(self):
        vid_lst = []
        prio_lst = []
        for pkt in self.__PKT_LST__:
            vid = str(pkt.vlan.id)
            vid_lst.append(vid)
            prio = str(pkt.vlan.priority)
            prio_lst.append(prio)

        return vid_lst, prio_lst

    def if_has_l2layer(self):
        try:
            if self.__PKT_LST__ is None:
                print('have l2 layer')
                return True
            else:
                print('2 no have l2 layer')
                return False
        except AttributeError as e:
            print('1no have l2 layer', e)
            return False

    def if_has_vid(self):
        try:
            for pkt in self.__PKT_LST__:
                if pkt.eth.vid:

                    print('have vlan layer')
                    return False
                else:
                    print('1no have vlan layer')
                    return False
        except AttributeError as e:
            print('no have vlan layer', e)
            return False

    def get_pkt_len(self):
        pkt_len = []
        for pkt in self.__PKT_LST__:
            pkt_lens = str(pkt.length)
            pkt_len.append(pkt_lens)
        return pkt_len

    def get_ip_dst(self):
        ip_da_lst = []
        for pkt in self.__PKT_LST__:
            ip_dst = str(pkt.ip.dst)
            ip_da_lst.append(ip_dst)
        return ip_da_lst

    def get_ip_src(self):
        ip_src_lst = []
        for pkt in self.__PKT_LST__:
            ip_src = str(pkt.ip.src)
            ip_src_lst.append(ip_src)
        return ip_src_lst

    def if_has_snmp(self):
        while True:
            for pkt in self.__PKT_LST__:
                if 'SNMP' in str(pkt.highest_layer):
                    print(pkt.highest_layer)
                    dst_port = str(pkt.udp.dstport)
                    src_port = str(pkt.udp.srcport)
                    # _snmp = pkt.snmp
                    # print(_snmp)
                    if '161' == dst_port or src_port == '161':
                        return True
            return False

if __name__ == '__main__':
    file1 = r"F:\Result\cap\test_0_1_vlan_untag.pcap"

    check = Check_pcap_res_base(file1)
    res, cnt = check.get_src_mac()
    for i in res:
        print(i)
        assert '00:00:20:00:00:00' == i
    vid, vcnt = check.if_has_vid()
    for i in vid:
        print(i)
